#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jun  8 15:01:31 2025

@author: fenghongli
"""

import os
import pandas as pd
from arch import arch_model

# 输入CSV目录
input_dir = 'stock_returns'
output_csv = 'garch_results.csv'
volatility_output_dir = 'garch_volatility'
os.makedirs(volatility_output_dir, exist_ok=True)

# 存储模型结果
results = []

for file in os.listdir(input_dir):
    if file.endswith('_returns.csv'):
        file_path = os.path.join(input_dir, file)
        try:
            df = pd.read_csv(file_path)
            ts_code = df['ts_code'].iloc[0]
            ret_series = df['ret'].dropna() * 100  # 将收益率转换为百分比更稳定

            # 拟合GARCH(1,1)模型
            model = arch_model(ret_series, vol='Garch', p=1, q=1, dist='normal')
            res = model.fit(disp='off')

            params = res.params
            pvalues = res.pvalues
            alpha = params['alpha[1]']
            beta = params['beta[1]']
            stability = '稳定' if (alpha + beta < 1) else '不稳定'

            results.append({
                'ts_code': ts_code,
                '文件名': file,
                'omega': round(params['omega'], 5),
                'alpha': round(alpha, 5),
                'beta': round(beta, 5),
                'alpha+beta': round(alpha + beta, 5),
                'alpha_p值': round(pvalues['alpha[1]'], 4),
                'beta_p值': round(pvalues['beta[1]'], 4),
                '稳定性判断': stability
            })

            # 保存条件波动率序列
            volatility_df = pd.DataFrame({
                'trade_date': df['trade_date'].iloc[-len(res.conditional_volatility):].values,
                'volatility': res.conditional_volatility
            })
            volatility_df.to_csv(
                os.path.join(volatility_output_dir, f"{ts_code}_volatility.csv"), index=False
            )

        except Exception as e:
            print(f"处理文件 {file} 出错：{e}")

# 保存模型参数结果
results_df = pd.DataFrame(results)
results_df.to_csv(output_csv, index=False, encoding='utf-8-sig')

print(f"GARCH模型拟合完成，参数保存在 {output_csv}，波动率序列保存至 {volatility_output_dir}/")
